import logging
import numpy as np
import random
from gym import spaces
import time
import gym
import random
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('TkAgg')


CHECKERBOARDSIZE = 12
BARRIERNUM = 5
REWARDNUM = 30
HIGHREWARDPRICE = 5
LOWREWARDPRICE = 1
# reward衰减
GAMMA = 1
BARRIERSTATE = -1


# 描述了棋盘的相关环境
class deliveryManEnv(gym.Env):
    logger = logging.getLogger(__name__)

    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 2
    }

    # 初始化一个n*n的棋盘
    def __init__(self):

        self.CHECKERBOARDSIZE = CHECKERBOARDSIZE
        self.states = np.zeros([CHECKERBOARDSIZE, CHECKERBOARDSIZE]) # 状态空间
        self.viewer = None

        # 给出随机障碍
        self.barrier_states = []
        while True:
            barrier_states_x = random.randint(0, CHECKERBOARDSIZE-1)
            barrier_states_y = random.randint(0, CHECKERBOARDSIZE-1)
            self.states[barrier_states_x, barrier_states_y] = -1
            self.barrier_states.append((barrier_states_x, barrier_states_y))
            if len(self.barrier_states) == BARRIERNUM:
                break

        # 给出随机的奖励
        # 给出指定的障碍
        self.reward_states = {}
        while True:
            randomReward_x = (random.randint(0, CHECKERBOARDSIZE-1))
            randomReward_y = (random.randint(0, CHECKERBOARDSIZE-1))
            if random in self.barrier_states:
                continue

            # 给出随机价格
            if(self.states[randomReward_x, randomReward_y]) == 0:
                random_reward = random.randint(LOWREWARDPRICE, HIGHREWARDPRICE)
                self.states[randomReward_x, randomReward_y] = random_reward
                self.reward_states[(randomReward_x, randomReward_y)] = random_reward
            if len(self.reward_states) == REWARDNUM:
                break

        # 给出随机初始
        initState_x = random.randint(0, CHECKERBOARDSIZE-1)
        initState_y = random.randint(0, CHECKERBOARDSIZE-1)
        if random not in self.barrier_states:
            self.states[initState_x, initState_y] = -2
            self.state = [initState_x, initState_y]

        self.init_states = self.states.copy()
        self.observation_space = self.states

    def step(self, action):
        is_terminal = False
        # 强制要求合法
        cur_state = self.state[:]
        if action == 'l':
            self.state[0] = self.state[0]-1
            self.states[cur_state[0], cur_state[1]] = 0
        elif action == 'r':
            self.state[0] = self.state[0]+1
            self.states[cur_state[0], cur_state[1]] = 0
        elif action == 'u':
            self.state[1] = self.state[1]-1
            self.states[cur_state[0], cur_state[1]] = 0
        elif action == 'd':
            self.state[1] = self.state[1]+1
            self.states[cur_state[0], cur_state[1]] = 0
        elif action == 'o':
            pass

        # 系统当前状态
        r = 0
        if self.states[self.state[0], self.state[1]] == -1:
            is_terminal = True
        elif self.states[self.state[0], self.state[1]] >= 0:
            r = self.states[self.state[0], self.state[1]]
            self.states[self.state[0], self.state[1]] = -2

        for i in range(CHECKERBOARDSIZE):
            for j in range(CHECKERBOARDSIZE):
                if self.states[i, j] > 0:
                    self.states[i, j] -= 1
        return self.states, r, is_terminal

    # 回复到本次最初的状态
    def reset(self):
        self.states = self.init_states.copy()
        return self.states

    def getTerminal(self):
        return self.terminate_states

    def getGamma(self):
        return self.gamma

    def getStates(self):
        return self.states

    def render(self, mode='human', close=False):
        bg_array = np.ones([12, 12]) * 100
        bg_array[self.state[0], self.state[1]] = -2
        fig, ax = plt.subplots()
        im = ax.imshow(bg_array)
        for i in range(12):
            for j in range(12):
                text = ax.text(j, i, self.states[i, j], ha="center", va="center", color='r')
        plt.show()
        plt.close()


    # 给出向量和方向转换的函数
    def actionToArray(self, direction):
        if direction == "e":
            return np.array([1, 0, 0, 0, 0])
        if direction == "w":
            return np.array([0, 1, 0, 0, 0])
        if direction == "s":
            return np.array([0, 0, 1, 0, 0])
        if direction == "n":
            return np.array([0, 0, 0, 1, 0])
        if direction == "o":
            return np.array([0, 0, 0, 0, 1])


    def arrayToAction(self, array):
        num = np.argmax(array)
        if num == 0:
            return "e"
        if num == 1:
            return "w"
        if num == 2:
            return "s"
        if num == 3:
            return "n"
        if num == 4:
            return "o"


env_dm = deliveryManEnv()

step_space = ['l', 'r', 'u', 'd', 'o']

for i in range(50):
    env_dm.step(random.choice(step_space))
    time.sleep(3)
    env_dm.render()



# a = env_dm.init_states
# fig, ax = plt.subplots()
# im = ax.imshow(a)
# for i in range(12):
#     for j in range(12):
#         text = ax.text(j, i, a[i, j], ha="center", va = "center", color = 'w')
# fig.tight_layout()
# plt.show()
#




